package com.yjw.zookeeper.distributed_lock;

import org.apache.zookeeper.*;
import org.apache.zookeeper.data.Stat;

import java.io.IOException;
import java.util.List;
import java.util.TreeSet;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.locks.Condition;
import java.util.concurrent.locks.Lock;

/**
 * 分布式锁实现
 */
public class DistributedLock implements Lock, Watcher {

    private ZooKeeper zooKeeper;
    // 定义根节点
    private String ROOT_NODE = "/locks";
    // 表示前一个节点
    private String PRE_NODE;
    // 表示当前节点
    private String CURRENT_NODE;
    private CountDownLatch countDownLatch;

    /**
     * 初始化数据
     */
    public DistributedLock() {
        try {
            zooKeeper = new ZooKeeper("192.168.202.18:2181,192.168.202.49:2181,192.168.202.50:2181",
                    4000, this);
            //等待连接进入CONNECTED状态
            int count = 0;
            while (true && count < 5) {
                if (zooKeeper.getState() == ZooKeeper.States.CONNECTED) {
                    break;
                }
                count++;
                Thread.sleep(500);
            }
            // 创建根节点
            Stat stat = zooKeeper.exists(ROOT_NODE, false);
            if (stat == null) {
                zooKeeper.create(ROOT_NODE, "0".getBytes(), ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
            }
        } catch (IOException e) {
            e.printStackTrace();
        } catch (InterruptedException e) {
            e.printStackTrace();
        } catch (KeeperException e) {
            e.printStackTrace();
        }
    }

    @Override
    public boolean tryLock() {
        try {
            // 创建临时有序节点
            CURRENT_NODE = zooKeeper.create(ROOT_NODE + "/", "0".getBytes(), ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.EPHEMERAL_SEQUENTIAL);
            System.out.println(Thread.currentThread().getName() + "->" + CURRENT_NODE + "，尝试竞争锁");
            // 获取根节点下所有子节点，并排序
            List<String> childrens = zooKeeper.getChildren(ROOT_NODE, false);
            TreeSet<String> treeSet = new TreeSet<>();
            for (String children : childrens) {
                treeSet.add(ROOT_NODE + "/" + children);
            }
            // 判断当前节点是不是最小节点
            String firstNode = treeSet.first();
            if (CURRENT_NODE.equals(firstNode)) {
                return true;
            }
            // 获取当前节点的前一个节点
            TreeSet<String> preTreeSet = (TreeSet<String>) treeSet.headSet(CURRENT_NODE);
            PRE_NODE = preTreeSet.last();
        } catch (KeeperException e) {
            e.printStackTrace();
        } catch (InterruptedException e) {
            e.printStackTrace();
        }
        return false;
    }

    @Override
    public void lock() {
        if (tryLock()) {
            System.out.println(Thread.currentThread().getName() + "->" + CURRENT_NODE + "，获得锁成功");
            return;
        }
        try {
            waitForLock(PRE_NODE);
        } catch (KeeperException e) {
            e.printStackTrace();
        } catch (InterruptedException e) {
            e.printStackTrace();
        }
    }

    /**
     * 等待锁
     *
     * @param preNode
     */
    private void waitForLock(String preNode) throws KeeperException, InterruptedException {
        // 监听当前节点的前一个节点
        Stat stat = zooKeeper.exists(preNode, true);
        if (stat != null) {
            System.out.println(Thread.currentThread().getName() + "->等待锁" + preNode + "释放");
            countDownLatch = new CountDownLatch(1);
            countDownLatch.await();
            System.out.println(Thread.currentThread().getName() + "->" + CURRENT_NODE + "，获得锁成功");
        }
        return;
    }

    @Override
    public void lockInterruptibly() throws InterruptedException {

    }

    @Override
    public boolean tryLock(long time, TimeUnit unit) throws InterruptedException {
        return false;
    }

    @Override
    public void unlock() {
        System.out.println(Thread.currentThread().getName() + "->释放锁" + CURRENT_NODE);
        try {
            zooKeeper.delete(CURRENT_NODE, -1);
            CURRENT_NODE = null;
            PRE_NODE = null;
            zooKeeper.close();
        } catch (InterruptedException e) {
            e.printStackTrace();
        } catch (KeeperException e) {
            e.printStackTrace();
        }

    }

    @Override
    public Condition newCondition() {
        return null;
    }

    @Override
    public void process(WatchedEvent event) {
        if (countDownLatch != null) {
            countDownLatch.countDown();
        }
    }
}
